import requests
from fastapi.responses import JSONResponse
from bs4 import BeautifulSoup
from urllib.parse import urljoin

class SiteImageAPI:
    def __init__(self):
        pass
        
    def handle(self, url: str):
        try:
            # 设置请求头，模拟浏览器访问
            headers = {
                'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36'
            }
            
            # 发送请求获取网页内容
            response = requests.get(url, headers=headers, timeout=10)
            response.raise_for_status()
            
            # 使用BeautifulSoup解析HTML
            soup = BeautifulSoup(response.text, 'html.parser')
            
            # 查找所有img标签
            img_tags = soup.find_all('img')
            
            # 提取图片URL并处理相对路径
            image_urls = {}
            for i, img in enumerate(img_tags):
                # 获取src属性
                src = img.get('src')
                if src:
                    # 处理相对URL，转换为绝对URL
                    absolute_url = urljoin(url, src)
                    image_urls[i+1] = absolute_url
            
            # return {
            #     "success": True,
            #     "message": f"成功获取网站图片，共{len(image_urls)}张",
            #     "data": {
            #         "image_count": len(image_urls),
            #         "images": image_urls
            #     }
            # }
            return JSONResponse(content={
                "success": True,
                "code": 200,
                "message": f"成功获取网站图片，共{len(image_urls)}张",
                "data": {
                    "image_count": len(image_urls),
                    "images": image_urls
                },
                "status_code": 200,
                "headers": response.headers
            })
            
        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "message": f"获取网站内容失败: {str(e)}",
                "data": None
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"处理网站内容时出错: {str(e)}",
                "data": None
            }